// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cstdint>
#include <memory>

#include <nebula/core/buffer.h>
#include <nebula/core/memory_pool.h>

#include <nebula/bits/bit_util.h>
#include <turbo/base/macros.h>

namespace nebula::internal {

    // A std::generate() like function to write sequential bits into a bitmap area.
    // Bits preceding the bitmap area are preserved, bits following the bitmap
    // area may be clobbered.

    template<class Generator>
    void GenerateBits(uint8_t *bitmap, int64_t start_offset, int64_t length, Generator &&g) {
        if (length == 0) {
            return;
        }
        uint8_t *cur = bitmap + start_offset / 8;
        uint8_t bit_mask = bit_util::kBitmask[start_offset % 8];
        uint8_t current_byte = *cur & bit_util::kPrecedingBitmask[start_offset % 8];

        for (int64_t index = 0; index < length; ++index) {
            const bool bit = g();
            current_byte = bit ? (current_byte | bit_mask) : current_byte;
            bit_mask = static_cast<uint8_t>(bit_mask << 1);
            if (bit_mask == 0) {
                bit_mask = 1;
                *cur++ = current_byte;
                current_byte = 0;
            }
        }
        if (bit_mask != 1) {
            *cur++ = current_byte;
        }
    }

    // Like GenerateBits(), but unrolls its main loop for higher performance.

    template<class Generator>
    void GenerateBitsUnrolled(uint8_t *bitmap, int64_t start_offset, int64_t length,
                              Generator &&g) {
        static_assert(std::is_same<decltype(std::declval<Generator>()()), bool>::value,
                      "Functor passed to GenerateBitsUnrolled must return bool");

        if (length == 0) {
            return;
        }
        uint8_t current_byte;
        uint8_t *cur = bitmap + start_offset / 8;
        const uint64_t start_bit_offset = start_offset % 8;
        uint8_t bit_mask = bit_util::kBitmask[start_bit_offset];
        int64_t remaining = length;

        if (bit_mask != 0x01) {
            current_byte = *cur & bit_util::kPrecedingBitmask[start_bit_offset];
            while (bit_mask != 0 && remaining > 0) {
                current_byte |= g() * bit_mask;
                bit_mask = static_cast<uint8_t>(bit_mask << 1);
                --remaining;
            }
            *cur++ = current_byte;
        }

        int64_t remaining_bytes = remaining / 8;
        uint8_t out_results[8];
        while (remaining_bytes-- > 0) {
            for (int i = 0; i < 8; ++i) {
                out_results[i] = g();
            }
            *cur++ = static_cast<uint8_t>(out_results[0] | out_results[1] << 1 |
                                          out_results[2] << 2 | out_results[3] << 3 |
                                          out_results[4] << 4 | out_results[5] << 5 |
                                          out_results[6] << 6 | out_results[7] << 7);
        }

        int64_t remaining_bits = remaining % 8;
        if (remaining_bits) {
            current_byte = 0;
            bit_mask = 0x01;
            while (remaining_bits-- > 0) {
                current_byte |= g() * bit_mask;
                bit_mask = static_cast<uint8_t>(bit_mask << 1);
            }
            *cur++ = current_byte;
        }
    }

}  // namespace nebula::internal
